package org.apache.giraph.benchmark;

import java.io.IOException;
import java.util.Map.Entry;

import org.apache.giraph.graph.BasicVertex;
import org.apache.giraph.graph.VertexWriter;
import org.apache.giraph.lib.TextVertexOutputFormat;
import org.apache.hadoop.io.FloatWritable;
import org.apache.hadoop.io.LongWritable;
import org.apache.hadoop.io.MapWritable;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;

public class DecisionTreeVertexOutputFormat extends
TextVertexOutputFormat<LongWritable, MapWritable, FloatWritable>{

	/**
	 * DecisionTreeVertexWriter
	 * @author genia
	 * 
	 * Write a decision tree to disk so that it does not have to be regenerated.
	 * Both RandomForestVertex and DecisionTreeVertex rely on this output format.
	 */
	public static class DecisionTreeVertexWriter extends
	TextVertexWriter<LongWritable, MapWritable, FloatWritable> {

		public DecisionTreeVertexWriter(
				RecordWriter<Text, Text> lineRecordWriter) {
			super(lineRecordWriter);
		}

		@Override
		public void writeVertex(
				BasicVertex<LongWritable, MapWritable, FloatWritable, ?> vertex)
		throws IOException, InterruptedException {

			// Only record current vertex if it is a tree node, vs a training/testing set node
			if (vertex.getVertexValue().get("type").toString() == "T") {
				
				// Initialize output string with vertex id
				// FIXME also record target attribute ID, once format is solidified
				String output = vertex.getVertexId().toString() + " - ";
				getRecordWriter().write(new Text(output), null);
				
				// Write child nodes and their classification decisions
				for (Entry<Writable, Writable> entry : 
					((MapWritable) vertex.getVertexValue().get("children")).entrySet()) {
					
					// Record attribute value followed by child value
					String child = entry.getKey().toString() + " : " + entry.getValue().toString();
					getRecordWriter().write(new Text(child), null);
					
				}
			}
		}
	}

	@Override
	public VertexWriter<LongWritable, MapWritable, FloatWritable> 
	createVertexWriter(TaskAttemptContext context) throws IOException,
			InterruptedException {
		RecordWriter<Text, Text> recordWriter = textOutputFormat.getRecordWriter(context);
		return new DecisionTreeVertexWriter(recordWriter);
	}

}
